from cmath import inf
from re import S
import re


import torch
import torch.nn as nn



class LinearClassifierProbe(nn.Module):
    def __init__(self, n_in_features, n_out_features, softmax_dim = None):
        super(LinearClassifierProbe, self).__init__()
        self.linear = nn.Linear(n_in_features, n_out_features)
        if softmax_dim:
            self.softmax = nn.Softmax(softmax_dim)

    def forward(self, in_features):
        out_features = self.linear(in_features)
        if hasattr(self, 'softmax'):
            out_features = self.softmax(out_features)
        return out_features

def valid_probe_and_get_loss(midlayers_features,probes,feature_name):
    with torch.no_grad():
        CEloss = nn.CrossEntropyLoss()
        X = torch.cat([midlayers_features[0][feature_name],midlayers_features[1][feature_name]],dim=0)
        preds = probes[feature_name](X)
        device = preds.device
        labels = torch.cat([torch.zeros(midlayers_features[0][feature_name].shape[0], dtype=torch.long), torch.ones(midlayers_features[1][feature_name].shape[0], dtype=torch.long)], 0).to(device)
        loss = CEloss(preds, labels)
    return loss.item()

def train_probe_and_get_loss(midlayers_features,probes,probe_optimizers,feature_name):
    # 开始训练
    CEloss = nn.CrossEntropyLoss()
    probe_optimizers[feature_name].zero_grad()
    # merge data
    X = torch.cat([midlayers_features[0][feature_name],midlayers_features[1][feature_name]],dim=0)
    # forward propagation
    preds = probes[feature_name](X)
    # labels
    device = preds.device
    labels = torch.cat([torch.zeros(midlayers_features[0][feature_name].shape[0], dtype=torch.long), torch.ones(midlayers_features[1][feature_name].shape[0], dtype=torch.long)], 0).to(device)
    # loss
    loss = CEloss(preds, labels)
    # backward propagation
    loss.backward()
    probe_optimizers[feature_name].step()
    return loss.item()

def get_probe_error(midlayers_features,probes,feature_name):
    pred_real = probes[feature_name](midlayers_features[0][feature_name])
    pred_gen = probes[feature_name](midlayers_features[1][feature_name])
    error = ((pred_real[:,0]<pred_real[:,1]).sum() + (pred_gen[:,0]>pred_gen[:,1]).sum()) \
        / ((pred_real.shape[0]+pred_gen.shape[0])*1.)
    return error.item()

if __name__ == "__main__":
    in_features = torch.randn(2, 10)
    print(in_features)
    n_in_features = in_features.shape[1]
    n_labels=2
    probe = LinearClassifierProbe(n_in_features, n_labels)
    print(probe(in_features))